from typing import Dict, Any, Optional, List, TypedDict, Annotated, Sequence, Set, Tuple, Callable, Union, TypeVar
from datetime import datetime
import asyncio
import json
import uuid
from dataclasses import dataclass, asdict, field
from enum import Enum, auto

from langgraph.graph import StateGraph, END, START
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode
from langgraph.checkpoint.memory import MemorySaver
# 使用MemorySaver替代SqliteSaver，简化依赖
# from langgraph.checkpoint.sqlite import SqliteSaver
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
from langchain_core.tools import tool

# 引入数据库 - 使用直接导入方式避免冲突
# 直接从database_core.py导入所需功能
from app.core.database_core import get_db, async_session_factory

# 重命名async_session_factory为async_session_maker以保持兼容性
async_session_maker = async_session_factory

# 导入async_db_pool模块
from app.core.database.async_db_pool import async_db_pool_manager, get_async_db_session
from app.core.db_utils import DatabaseUtils

# 修复logger导入
from loguru import logger as _logger
def get_logger(name):
    return _logger.bind(component=name)
from utils.exceptions import ToolError, WorkflowError

logger = get_logger("langgraph_workflow")


class AgentStatus(Enum):
    """智能体状态"""
    IDLE = "idle"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"
    TIMEOUT = "timeout"
    RETRYING = "retrying"
    PAUSED = "paused"


class TaskType(Enum):
    """任务类型"""
    TEXT_RETRIEVAL = "text_retrieval"
    DOC_SEGMENT = "doc_segment"
    MULTI_MODAL = "multi_modal"
    VALIDATION = "validation"
    INTEGRATION = "integration"
    FALLBACK = "fallback"
    DATA_PROCESSING = "data_processing"
    KNOWLEDGE_UPDATE = "knowledge_update"
    METADATA_EXTRACTION = "metadata_extraction"


class AgentState(TypedDict):
    """智能体状态"""
    messages: Annotated[Sequence[BaseMessage], add_messages]
    task_type: str
    query: str
    context: Dict[str, Any]
    agent_status: Dict[str, AgentStatus]
    results: Dict[str, Any]
    errors: List[str]
    metadata: Dict[str, Any]
    timestamp: datetime
    # 新增字段用于任务依赖管理
    task_dependencies: Dict[str, List[str]]  # 任务ID -> 依赖的任务ID列表
    completed_tasks: Set[str]  # 已完成的任务ID集合
    current_task_id: Optional[str]  # 当前执行的任务ID
    loop_count: Dict[str, int]  # 循环计数
    max_iterations: Dict[str, int]  # 最大迭代次数
    workflow_id: str  # 工作流实例ID


@dataclass
class WorkflowConfig:
    """工作流配置"""
    max_retries: int = 3
    timeout_seconds: int = 30
    enable_validation: bool = True
    enable_fallback: bool = True
    parallel_execution: bool = True
    result_aggregation: str = "weighted"  # weighted, majority, priority
    checkpoint_type: str = "sqlite"  # memory, sqlite, database
    checkpoint_path: str = "./workflow_checkpoints.db"
    enable_persistence: bool = True
    enable_task_dependencies: bool = True
    enable_visualization: bool = False
    max_parallel_tasks: int = 5
    default_max_iterations: int = 10


class LangGraphWorkflowEngine:
    """LangGraph工作流引擎"""

    def __init__(self, config: Optional[WorkflowConfig] = None):
        self.config = config or WorkflowConfig()
        self.graph = None
        self.memory = None
        self.agents = {}
        self.tools = {}
        self.task_definitions = {}
        self.task_dependencies = {}
        self.db_utils = None
        self._init_checkpoint()
        self._build_graph()
        self._init_database()
    
    def _init_checkpoint(self):
        """初始化检查点存储"""
        if not self.config.enable_persistence:
            self.memory = MemorySaver()
        elif self.config.checkpoint_type == "sqlite":
            # 使用MemorySaver替代SqliteSaver
            self.memory = MemorySaver()
        elif self.config.checkpoint_type == "database":
            # 后续可以实现基于PostgreSQL的持久化
            self.memory = MemorySaver()
            logger.warning("基于数据库的检查点存储尚未实现，使用内存存储作为回退")
        else:
            self.memory = MemorySaver()
        logger.info(f"检查点存储初始化完成: {self.config.checkpoint_type}")
    
    def _init_database(self):
        """初始化数据库工具"""
        try:
            self.db_utils = DatabaseUtils()
            logger.info("数据库工具初始化完成")
        except Exception as e:
            logger.error(f"数据库工具初始化失败: {str(e)}")

    def register_agent(self, agent_name: str, agent_instance):
        """注册智能体"""
        self.agents[agent_name] = agent_instance
        logger.info(f"智能体注册完成: {agent_name}")

    def register_tool(self, tool_name: str, tool_function):
        """注册工具"""
        self.tools[tool_name] = tool(tool_function)
        logger.info(f"工具注册完成: {tool_name}")
    
    def register_task(self, task_id: str, task_handler: Callable, dependencies: Optional[List[str]] = None):
        """注册任务及其依赖关系
        
        Args:
            task_id: 任务ID
            task_handler: 任务处理函数
            dependencies: 依赖的任务ID列表
        """
        self.task_definitions[task_id] = task_handler
        if dependencies:
            self.task_dependencies[task_id] = dependencies
        logger.info(f"任务注册完成: {task_id}, 依赖: {dependencies}")
    
    def add_task_dependency(self, task_id: str, depends_on: str):
        """添加任务依赖关系
        
        Args:
            task_id: 任务ID
            depends_on: 依赖的任务ID
        """
        if task_id not in self.task_dependencies:
            self.task_dependencies[task_id] = []
        if depends_on not in self.task_dependencies[task_id]:
            self.task_dependencies[task_id].append(depends_on)
        logger.info(f"添加任务依赖: {task_id} -> {depends_on}")
    
    def _check_task_dependencies(self, task_id: str, completed_tasks: Set[str]) -> bool:
        """检查任务依赖是否满足
        
        Args:
            task_id: 任务ID
            completed_tasks: 已完成的任务集合
            
        Returns:
            依赖是否满足
        """
        if task_id not in self.task_dependencies:
            return True
        
        dependencies = self.task_dependencies[task_id]
        return all(dep in completed_tasks for dep in dependencies)

    def _build_graph(self):
        """构建工作流图"""
        # 创建工作流图
        workflow = StateGraph(AgentState)
        
        # 添加核心节点
        workflow.add_node("initialize", self._initialize_node)
        workflow.add_node("router", self._router_node)
        workflow.add_node("text_retrieval", self._text_retrieval_node)
        workflow.add_node("doc_segment", self._doc_segment_node)
        workflow.add_node("multi_modal", self._multi_modal_node)
        workflow.add_node("validation", self._validation_node)
        workflow.add_node("integration", self._integration_node)
        workflow.add_node("fallback", self._fallback_node)
        workflow.add_node("task_manager", self._task_manager_node)
        workflow.add_node("loop_handler", self._loop_handler_node)
        workflow.add_node("persist_state", self._persist_state_node)
        
        # 设置入口点
        workflow.set_entry_point("initialize")
        
        # 初始化 -> 路由
        workflow.add_edge("initialize", "router")
        
        # 路由 -> 功能节点
        workflow.add_edge("router", "text_retrieval")
        workflow.add_edge("router", "doc_segment")
        workflow.add_edge("router", "multi_modal")
        
        # 功能节点 -> 条件分支
        for node in ["text_retrieval", "doc_segment", "multi_modal"]:
            workflow.add_conditional_edges(
                node,
                self._should_validate,
                {
                    "validate": "validation",
                    "skip": "task_manager"
                }
            )
        
        # 验证 -> 任务管理器
        workflow.add_edge("validation", "task_manager")
        
        # 任务管理器 -> 条件分支（循环处理或继续）
        workflow.add_conditional_edges(
            "task_manager",
            self._should_loop,
            {
                "loop": "loop_handler",
                "continue": "integration"
            }
        )
        
        # 循环处理器 -> 路由（重新路由）
        workflow.add_edge("loop_handler", "router")
        
        # 整合 -> 状态持久化
        workflow.add_edge("integration", "persist_state")
        
        # 状态持久化 -> 条件分支（回退或结束）
        workflow.add_conditional_edges(
            "persist_state",
            self._should_fallback,
            {
                "fallback": "fallback",
                "end": END
            }
        )
        
        # 回退 -> 结束
        workflow.add_edge("fallback", END)
        
        # 编译图
        self.graph = workflow.compile(checkpointer=self.memory)
        logger.info("工作流图构建完成")
        
        # 注册预定义任务
        self._register_predefined_tasks()
    
    def _register_predefined_tasks(self):
        """注册预定义任务"""
        # 注册核心任务及其依赖关系
        self.register_task("text_retrieval", self._text_retrieval_node)
        self.register_task("doc_segment", self._doc_segment_node)
        self.register_task("multi_modal", self._multi_modal_node)
        self.register_task("validation", self._validation_node, ["text_retrieval", "doc_segment", "multi_modal"])
        self.register_task("integration", self._integration_node, ["validation"])
        self.register_task("persist_state", self._persist_state_node, ["integration"])

    async def _initialize_node(self, state: AgentState) -> Dict[str, Any]:
        """初始化节点"""
        try:
            logger.info("工作流初始化")
            
            # 初始化工作流ID
            workflow_id = state.get("workflow_id", str(uuid.uuid4()))
            
            # 初始化任务依赖相关字段
            state["workflow_id"] = workflow_id
            state["task_dependencies"] = self.task_dependencies.copy()
            state["completed_tasks"] = set()
            state["loop_count"] = {}
            state["max_iterations"] = {}
            
            # 初始化所有智能体状态
            for agent in ["initialize", "router", "text_retrieval", "doc_segment", "multi_modal",
                         "validation", "integration", "fallback", "task_manager", "loop_handler", "persist_state"]:
                if agent not in state["agent_status"]:
                    state["agent_status"][agent] = AgentStatus.IDLE
            
            state["agent_status"]["initialize"] = AgentStatus.COMPLETED
            state["timestamp"] = datetime.now()
            
            logger.info(f"工作流初始化完成: workflow_id={workflow_id}")
            return {"messages": [AIMessage(content=f"工作流初始化完成: {workflow_id}")]}
            
        except Exception as e:
            logger.error(f"初始化节点失败: {str(e)}")
            state["errors"].append(f"初始化失败: {str(e)}")
            state["agent_status"]["initialize"] = AgentStatus.FAILED
            return {"messages": [AIMessage(content=f"初始化失败: {str(e)}")]}
    
    async def _router_node(self, state: AgentState) -> Dict[str, Any]:
        """路由节点"""
        try:
            query = state["query"]
            task_type = state["task_type"]
            current_task_id = state.get("current_task_id")
            
            logger.info(f"路由节点处理: query='{query}', task_type='{task_type}', current_task='{current_task_id}'")
            
            # 更新智能体状态
            state["agent_status"]["router"] = AgentStatus.PROCESSING
            state["timestamp"] = datetime.now()
            
            # 如果有当前任务且依赖检查开启，检查依赖
            if current_task_id and self.config.enable_task_dependencies:
                if not self._check_task_dependencies(current_task_id, state["completed_tasks"]):
                    dependencies = self.task_dependencies.get(current_task_id, [])
                    missing = [dep for dep in dependencies if dep not in state["completed_tasks"]]
                    logger.warning(f"任务依赖未满足: {current_task_id} 需要 {missing}")
                    return {"messages": [AIMessage(content=f"任务依赖未满足，等待前置任务完成")]}
            
            # 根据任务类型路由到不同的智能体
            if task_type == TaskType.TEXT_RETRIEVAL.value:
                state["current_task_id"] = "text_retrieval"
                return {"messages": [AIMessage(content="路由到文本检索智能体")]}
            elif task_type == TaskType.DOC_SEGMENT.value:
                state["current_task_id"] = "doc_segment"
                return {"messages": [AIMessage(content="路由到文档分段检索智能体")]}
            elif task_type == TaskType.MULTI_MODAL.value:
                state["current_task_id"] = "multi_modal"
                return {"messages": [AIMessage(content="路由到多模态智能体")]}
            elif task_type == TaskType.DATA_PROCESSING.value:
                state["current_task_id"] = "data_processing"
                return {"messages": [AIMessage(content="路由到数据处理智能体")]}
            elif task_type == TaskType.KNOWLEDGE_UPDATE.value:
                state["current_task_id"] = "knowledge_update"
                return {"messages": [AIMessage(content="路由到知识更新智能体")]}
            else:
                return {"messages": [AIMessage(content="路由到默认智能体")]}
                
        except Exception as e:
            logger.error(f"路由节点失败: {str(e)}")
            state["errors"].append(f"路由失败: {str(e)}")
            state["agent_status"]["router"] = AgentStatus.FAILED
            return {"messages": [AIMessage(content=f"路由失败: {str(e)}")]}

    async def _text_retrieval_node(self, state: AgentState) -> Dict[str, Any]:
        """文本检索节点"""
        try:
            query = state["query"]
            context = state["context"]
            
            logger.info(f"文本检索节点处理: query='{query}'")
            
            state["agent_status"]["text_retrieval"] = AgentStatus.PROCESSING
            
            # 重试机制
            retries = 0
            while retries < self.config.max_retries:
                try:
                    # 调用文本检索智能体
                    text_agent = self.agents.get("text_qa")
                    if text_agent:
                        result = await text_agent.process_task(query, context)
                        state["results"]["text_retrieval"] = result
                        state["completed_tasks"].add("text_retrieval")
                        state["agent_status"]["text_retrieval"] = AgentStatus.COMPLETED
                        
                        # 更新循环计数
                        state["loop_count"]["text_retrieval"] = state["loop_count"].get("text_retrieval", 0) + 1
                        
                        return {"messages": [AIMessage(content=f"文本检索完成: {str(result)}")]}
                    else:
                        raise ToolError("文本检索智能体未注册")
                except Exception as e:
                    retries += 1
                    if retries >= self.config.max_retries:
                        raise
                    state["agent_status"]["text_retrieval"] = AgentStatus.RETRYING
                    logger.warning(f"文本检索重试 {retries}/{self.config.max_retries}: {str(e)}")
                    await asyncio.sleep(1)  # 短暂延迟后重试
                
        except Exception as e:
            logger.error(f"文本检索节点失败: {str(e)}")
            state["errors"].append(f"文本检索失败: {str(e)}")
            state["agent_status"]["text_retrieval"] = AgentStatus.FAILED
            return {"messages": [AIMessage(content=f"文本检索失败: {str(e)}")]}

    async def _doc_segment_node(self, state: AgentState) -> Dict[str, Any]:
        """文档分段检索节点"""
        try:
            query = state["query"]
            context = state["context"]
            
            logger.info(f"文档分段检索节点处理: query='{query}'")
            
            state["agent_status"]["doc_segment"] = AgentStatus.PROCESSING
            
            # 重试机制
            retries = 0
            while retries < self.config.max_retries:
                try:
                    # 调用文档分段检索智能体
                    doc_agent = self.agents.get("doc_segment")
                    if doc_agent:
                        result = await doc_agent.process_task(query, context)
                        state["results"]["doc_segment"] = result
                        state["completed_tasks"].add("doc_segment")
                        state["agent_status"]["doc_segment"] = AgentStatus.COMPLETED
                        
                        # 更新循环计数
                        state["loop_count"]["doc_segment"] = state["loop_count"].get("doc_segment", 0) + 1
                        
                        return {"messages": [AIMessage(content=f"文档分段检索完成: {str(result)}")]}
                    else:
                        raise ToolError("文档分段检索智能体未注册")
                except Exception as e:
                    retries += 1
                    if retries >= self.config.max_retries:
                        raise
                    state["agent_status"]["doc_segment"] = AgentStatus.RETRYING
                    logger.warning(f"文档分段检索重试 {retries}/{self.config.max_retries}: {str(e)}")
                    await asyncio.sleep(1)  # 短暂延迟后重试
                
        except Exception as e:
            logger.error(f"文档分段检索节点失败: {str(e)}")
            state["errors"].append(f"文档分段检索失败: {str(e)}")
            state["agent_status"]["doc_segment"] = AgentStatus.FAILED
            return {"messages": [AIMessage(content=f"文档分段检索失败: {str(e)}")]}

    async def _multi_modal_node(self, state: AgentState) -> Dict[str, Any]:
        """多模态节点"""
        try:
            query = state["query"]
            context = state["context"]
            
            logger.info(f"多模态节点处理: query='{query}'")
            
            state["agent_status"]["multi_modal"] = AgentStatus.PROCESSING
            
            # 重试机制
            retries = 0
            while retries < self.config.max_retries:
                try:
                    # 调用多模态智能体
                    multi_modal_agent = self.agents.get("multi_modal")
                    if multi_modal_agent:
                        result = await multi_modal_agent.process_task(query, context)
                        state["results"]["multi_modal"] = result
                        state["completed_tasks"].add("multi_modal")
                        state["agent_status"]["multi_modal"] = AgentStatus.COMPLETED
                        
                        # 更新循环计数
                        state["loop_count"]["multi_modal"] = state["loop_count"].get("multi_modal", 0) + 1
                        
                        return {"messages": [AIMessage(content=f"多模态处理完成: {str(result)}")]}
                    else:
                        raise ToolError("多模态智能体未注册")
                except Exception as e:
                    retries += 1
                    if retries >= self.config.max_retries:
                        raise
                    state["agent_status"]["multi_modal"] = AgentStatus.RETRYING
                    logger.warning(f"多模态处理重试 {retries}/{self.config.max_retries}: {str(e)}")
                    await asyncio.sleep(1)  # 短暂延迟后重试
                
        except Exception as e:
            logger.error(f"多模态节点失败: {str(e)}")
            state["errors"].append(f"多模态处理失败: {str(e)}")
            state["agent_status"]["multi_modal"] = AgentStatus.FAILED
            return {"messages": [AIMessage(content=f"多模态处理失败: {str(e)}")]}

    async def _validation_node(self, state: AgentState) -> Dict[str, Any]:
        """验证节点"""
        try:
            logger.info("验证节点处理")
            
            state["agent_status"]["validation"] = AgentStatus.PROCESSING
            
            # 检查是否满足依赖条件
            if not self._check_task_dependencies("validation", state["completed_tasks"]):
                missing_deps = [dep for dep in self.task_dependencies.get("validation", []) 
                               if dep not in state["completed_tasks"]]
                logger.warning(f"验证任务依赖未满足: 缺少 {missing_deps}")
                state["agent_status"]["validation"] = AgentStatus.PAUSED
                return {"messages": [AIMessage(content=f"验证任务暂停: 缺少依赖 {missing_deps}")]}
            
            # 调用单点校验智能体
            validation_agent = self.agents.get("validation")
            if validation_agent:
                results = state["results"]
                validation_result = await validation_agent.validate_results(results)
                
                state["results"]["validation"] = validation_result
                state["completed_tasks"].add("validation")
                state["agent_status"]["validation"] = AgentStatus.COMPLETED
                
                # 记录验证元数据
                state["metadata"]["validation"] = {
                    "timestamp": datetime.now().isoformat(),
                    "is_valid": validation_result.get("is_valid", False),
                    "errors": validation_result.get("errors", [])
                }
                
                if validation_result.get("is_valid", False):
                    return {"messages": [AIMessage(content="验证通过")]}
                else:
                    state["errors"].append(f"验证失败: {validation_result.get('errors', [])}")
                    return {"messages": [AIMessage(content=f"验证失败: {validation_result.get('errors', [])}")]}
            else:
                # 如果没有验证智能体，直接通过
                state["completed_tasks"].add("validation")
                state["agent_status"]["validation"] = AgentStatus.COMPLETED
                return {"messages": [AIMessage(content="跳过验证")]}
                
        except Exception as e:
            logger.error(f"验证节点失败: {str(e)}")
            state["errors"].append(f"验证失败: {str(e)}")
            state["agent_status"]["validation"] = AgentStatus.FAILED
            return {"messages": [AIMessage(content=f"验证失败: {str(e)}")]}

    async def _integration_node(self, state: AgentState) -> Dict[str, Any]:
        """整合节点"""
        try:
            logger.info("整合节点处理")
            
            state["agent_status"]["integration"] = AgentStatus.PROCESSING
            
            # 检查是否满足依赖条件
            if not self._check_task_dependencies("integration", state["completed_tasks"]):
                missing_deps = [dep for dep in self.task_dependencies.get("integration", []) 
                               if dep not in state["completed_tasks"]]
                logger.warning(f"整合任务依赖未满足: 缺少 {missing_deps}")
                state["agent_status"]["integration"] = AgentStatus.PAUSED
                return {"messages": [AIMessage(content=f"整合任务暂停: 缺少依赖 {missing_deps}")]}
            
            # 调用整合智能体
            integration_agent = self.agents.get("integration")
            if integration_agent:
                results = state["results"]
                integration_result = await integration_agent.integrate_results(results)
                
                state["results"]["integration"] = integration_result
                state["completed_tasks"].add("integration")
                state["agent_status"]["integration"] = AgentStatus.COMPLETED
                
                return {"messages": [AIMessage(content=f"整合完成: {str(integration_result)}")]}
            else:
                # 如果没有整合智能体，简单合并结果
                final_result = self._simple_integration(state["results"])
                state["results"]["integration"] = final_result
                state["completed_tasks"].add("integration")
                state["agent_status"]["integration"] = AgentStatus.COMPLETED
                return {"messages": [AIMessage(content=f"简单整合完成: {str(final_result)}")]}
                
        except Exception as e:
            logger.error(f"整合节点失败: {str(e)}")
            state["errors"].append(f"整合失败: {str(e)}")
            state["agent_status"]["integration"] = AgentStatus.FAILED
            return {"messages": [AIMessage(content=f"整合失败: {str(e)}")]}
    
    async def _task_manager_node(self, state: AgentState) -> Dict[str, Any]:
        """任务管理节点"""
        try:
            logger.info("任务管理节点处理")
            
            state["agent_status"]["task_manager"] = AgentStatus.PROCESSING
            
            # 检查任务完成情况
            completed_count = len(state["completed_tasks"])
            total_tasks = len(self.task_definitions)
            
            logger.info(f"任务完成情况: {completed_count}/{total_tasks}")
            
            # 检查是否有失败的任务需要重试
            failed_tasks = [agent for agent, status in state["agent_status"].items() 
                           if status == AgentStatus.FAILED]
            
            if failed_tasks:
                logger.warning(f"发现失败任务: {failed_tasks}")
                state["metadata"]["failed_tasks"] = failed_tasks
            
            state["agent_status"]["task_manager"] = AgentStatus.COMPLETED
            
            return {"messages": [AIMessage(content=f"任务管理: {completed_count}/{total_tasks} 任务完成")]}
            
        except Exception as e:
            logger.error(f"任务管理节点失败: {str(e)}")
            state["errors"].append(f"任务管理失败: {str(e)}")
            state["agent_status"]["task_manager"] = AgentStatus.FAILED
            return {"messages": [AIMessage(content=f"任务管理失败: {str(e)}")]}
    
    async def _loop_handler_node(self, state: AgentState) -> Dict[str, Any]:
        """循环处理节点"""
        try:
            logger.info("循环处理节点处理")
            
            state["agent_status"]["loop_handler"] = AgentStatus.PROCESSING
            
            # 更新循环状态
            current_task = state.get("current_task_id")
            if current_task:
                # 增加循环计数
                state["loop_count"][current_task] = state["loop_count"].get(current_task, 0) + 1
                loop_count = state["loop_count"][current_task]
                
                # 获取最大迭代次数
                max_iterations = state["max_iterations"].get(
                    current_task, 
                    self.config.default_max_iterations
                )
                
                logger.info(f"循环处理: 任务={current_task}, 次数={loop_count}/{max_iterations}")
                
                # 检查是否超过最大迭代次数
                if loop_count >= max_iterations:
                    logger.warning(f"任务 {current_task} 达到最大迭代次数 {max_iterations}")
                    state["errors"].append(f"任务 {current_task} 达到最大迭代次数")
            
            state["agent_status"]["loop_handler"] = AgentStatus.COMPLETED
            
            return {"messages": [AIMessage(content="循环处理完成")]}
            
        except Exception as e:
            logger.error(f"循环处理节点失败: {str(e)}")
            state["errors"].append(f"循环处理失败: {str(e)}")
            state["agent_status"]["loop_handler"] = AgentStatus.FAILED
            return {"messages": [AIMessage(content=f"循环处理失败: {str(e)}")]}
    
    async def _persist_state_node(self, state: AgentState) -> Dict[str, Any]:
        """状态持久化节点"""
        try:
            logger.info("状态持久化节点处理")
            
            state["agent_status"]["persist_state"] = AgentStatus.PROCESSING
            
            # 检查是否满足依赖条件
            if not self._check_task_dependencies("persist_state", state["completed_tasks"]):
                missing_deps = [dep for dep in self.task_dependencies.get("persist_state", []) 
                               if dep not in state["completed_tasks"]]
                logger.warning(f"持久化任务依赖未满足: 缺少 {missing_deps}")
                state["agent_status"]["persist_state"] = AgentStatus.PAUSED
                return {"messages": [AIMessage(content=f"持久化任务暂停: 缺少依赖 {missing_deps}")]}
            
            # 使用数据库工具持久化状态
            if self.db_utils:
                try:
                    # 构建要保存的数据
                    workflow_data = {
                        "workflow_id": state["workflow_id"],
                        "task_type": state["task_type"],
                        "query": state["query"],
                        "results": state["results"],
                        "errors": state["errors"],
                        "agent_status": {k: v.value for k, v in state["agent_status"].items()},
                        "completed_tasks": list(state["completed_tasks"]),
                        "metadata": state["metadata"],
                        "timestamp": state["timestamp"].isoformat()
                    }
                    
                    # 保存到数据库
                    await self.db_utils.save_workflow_state(workflow_data)
                    logger.info(f"工作流状态持久化成功: {state['workflow_id']}")
                    
                except Exception as db_error:
                    logger.error(f"数据库持久化失败: {str(db_error)}")
                    # 不抛出异常，继续执行
            
            state["completed_tasks"].add("persist_state")
            state["agent_status"]["persist_state"] = AgentStatus.COMPLETED
            
            return {"messages": [AIMessage(content="状态持久化完成")]}
            
        except Exception as e:
            logger.error(f"状态持久化节点失败: {str(e)}")
            state["errors"].append(f"状态持久化失败: {str(e)}")
            state["agent_status"]["persist_state"] = AgentStatus.FAILED
            return {"messages": [AIMessage(content=f"状态持久化失败: {str(e)}")]}

    async def _fallback_node(self, state: AgentState) -> Dict[str, Any]:
        """回退节点"""
        try:
            logger.info("回退节点处理")
            
            state["agent_status"]["fallback"] = AgentStatus.PROCESSING
            
            # 调用回退智能体
            fallback_agent = self.agents.get("fallback")
            if fallback_agent:
                errors = state["errors"]
                fallback_result = await fallback_agent.handle_fallback(errors)
                
                state["results"]["fallback"] = fallback_result
                state["agent_status"]["fallback"] = AgentStatus.COMPLETED
                
                return {"messages": [AIMessage(content=f"回退处理完成: {str(fallback_result)}")]}
            else:
                # 如果没有回退智能体，提供默认回退
                default_fallback = self._default_fallback(state["errors"])
                state["results"]["fallback"] = default_fallback
                state["agent_status"]["fallback"] = AgentStatus.COMPLETED
                return {"messages": [AIMessage(content=f"默认回退完成: {str(default_fallback)}")]}
                
        except Exception as e:
            logger.error(f"回退节点失败: {str(e)}")
            state["errors"].append(f"回退处理失败: {str(e)}")
            state["agent_status"]["fallback"] = AgentStatus.FAILED
            return {"messages": [AIMessage(content=f"回退处理失败: {str(e)}")]}

    def _should_validate(self, state: AgentState) -> str:
        """判断是否需要验证"""
        if self.config.enable_validation:
            return "validate"
        else:
            return "skip"

    def _should_fallback(self, state: AgentState) -> str:
        """判断是否需要回退"""
        if state["errors"] and self.config.enable_fallback:
            # 检查是否有严重错误
            has_critical_error = any("critical" in error.lower() for error in state["errors"])
            # 检查失败任务数量
            failed_tasks = [agent for agent, status in state["agent_status"].items() 
                           if status == AgentStatus.FAILED]
            
            if has_critical_error or len(failed_tasks) > 1:
                return "fallback"
        
        # 即使有错误，如果启用了任务依赖，也可能需要继续执行
        if self.config.enable_task_dependencies:
            # 检查是否还有未完成但可执行的任务
            for task_id in self.task_definitions:
                if task_id not in state["completed_tasks"] and \
                   task_id not in [agent for agent, status in state["agent_status"].items() if status == AgentStatus.FAILED] and \
                   self._check_task_dependencies(task_id, state["completed_tasks"]):
                    return "continue"
        
        return "end"
    
    def _should_loop(self, state: AgentState) -> str:
        """判断是否需要循环执行"""
        # 检查是否启用了循环
        if "loop_enabled" in state.get("metadata", {}) and state["metadata"]["loop_enabled"]:
            current_task = state.get("current_task_id")
            if current_task:
                loop_count = state["loop_count"].get(current_task, 0)
                max_iterations = state["max_iterations"].get(
                    current_task, 
                    self.config.default_max_iterations
                )
                
                # 检查是否达到最大迭代次数
                if loop_count < max_iterations:
                    # 检查是否需要继续循环的条件
                    if self._should_continue_loop(state):
                        return "loop"
        
        return "continue"
    
    def _should_continue_loop(self, state: AgentState) -> bool:
        """判断是否应该继续循环"""
        # 这里可以实现更复杂的循环条件判断逻辑
        # 例如：检查结果质量、用户反馈等
        
        # 简单实现：如果没有错误且结果不完整，则继续循环
        has_errors = len(state["errors"]) > 0
        
        # 检查核心任务是否都已完成
        core_tasks = ["text_retrieval", "doc_segment", "multi_modal"]
        all_core_completed = all(task in state["completed_tasks"] for task in core_tasks if task in self.task_definitions)
        
        # 如果没有错误但核心任务未完成，则继续循环
        if not has_errors and not all_core_completed:
            return True
        
        return False

    def _simple_integration(self, results: Dict[str, Any]) -> Dict[str, Any]:
        """简单整合结果"""
        # 提取所有成功的结果
        successful_results = {}
        for agent_name, result in results.items():
            if isinstance(result, dict) and result.get("success", False):
                successful_results[agent_name] = result
        
        # 根据配置选择整合策略
        if self.config.result_aggregation == "weighted":
            return self._weighted_aggregation(successful_results)
        elif self.config.result_aggregation == "majority":
            return self._majority_voting(successful_results)
        else:
            return self._priority_based(successful_results)

    def _weighted_aggregation(self, results: Dict[str, Any]) -> Dict[str, Any]:
        """加权聚合"""
        # 定义权重
        weights = {
            "text_retrieval": 0.3,
            "doc_segment": 0.4,
            "multi_modal": 0.3
        }
        
        final_result = {
            "success": True,
            "method": "weighted_aggregation",
            "results": results,
            "weights": weights,
            "confidence": 0.0
        }
        
        # 计算置信度
        total_weight = 0
        for agent, result in results.items():
            weight = weights.get(agent, 0.1)
            confidence = result.get("confidence", 0.5)
            total_weight += weight * confidence
        
        final_result["confidence"] = min(total_weight, 1.0)
        return final_result

    def _majority_voting(self, results: Dict[str, Any]) -> Dict[str, Any]:
        """多数投票"""
        # 这里可以实现基于答案的多数投票
        return {
            "success": True,
            "method": "majority_voting",
            "results": results,
            "confidence": len(results) / 3.0  # 假设有3个智能体
        }

    def _priority_based(self, results: Dict[str, Any]) -> Dict[str, Any]:
        """基于优先级"""
        # 按优先级选择结果
        priority_order = ["doc_segment", "text_retrieval", "multi_modal"]
        
        for priority in priority_order:
            if priority in results:
                return {
                    "success": True,
                    "method": "priority_based",
                    "selected_agent": priority,
                    "result": results[priority],
                    "confidence": results[priority].get("confidence", 0.5)
                }
        
        return {
            "success": False,
            "method": "priority_based",
            "error": "没有可用结果"
        }

    def _default_fallback(self, errors: List[str]) -> Dict[str, Any]:
        """默认回退处理"""
        return {
            "success": False,
            "method": "default_fallback",
            "errors": errors,
            "fallback_message": "系统遇到错误，请稍后重试或联系技术支持。",
            "suggestions": [
                "检查输入参数是否正确",
                "确认文档格式是否支持",
                "尝试简化查询语句",
                "检查网络连接状态"
            ]
        }

    async def execute_workflow(self, query: str, task_type: str, 
                              context: Optional[Dict[str, Any]] = None,
                              thread_id: Optional[str] = None,
                              workflow_id: Optional[str] = None,
                              enable_loop: bool = False,
                              max_iterations: Optional[int] = None) -> Dict[str, Any]:
        """执行工作流
        
        Args:
            query: 查询内容
            task_type: 任务类型
            context: 上下文信息
            thread_id: 线程ID（用于状态持久化）
            workflow_id: 工作流实例ID
            enable_loop: 是否启用循环执行
            max_iterations: 最大迭代次数
            
        Returns:
            工作流执行结果
        """
        try:
            # 生成工作流ID
            actual_workflow_id = workflow_id or str(uuid.uuid4())
            actual_thread_id = thread_id or actual_workflow_id
            
            # 初始化状态
            initial_state = AgentState(
                messages=[HumanMessage(content=query)],
                task_type=task_type,
                query=query,
                context=context or {},
                agent_status={agent: AgentStatus.IDLE for agent in [
                    "initialize", "router", "text_retrieval", "doc_segment", "multi_modal",
                    "validation", "integration", "fallback", "task_manager", "loop_handler", "persist_state"
                ]},
                results={},
                errors=[],
                metadata={
                    "loop_enabled": enable_loop,
                    "max_iterations": max_iterations or self.config.default_max_iterations,
                    "workflow_start_time": datetime.now().isoformat()
                },
                timestamp=datetime.now(),
                task_dependencies={},
                completed_tasks=set(),
                current_task_id=None,
                loop_count={},
                max_iterations={},
                workflow_id=actual_workflow_id
            )
            
            # 设置最大迭代次数
            if max_iterations:
                initial_state["max_iterations"]["default"] = max_iterations
            
            # 配置线程
            config = {"configurable": {"thread_id": actual_thread_id}}
            
            logger.info(f"开始执行工作流: workflow_id='{actual_workflow_id}', task_type='{task_type}', query='{query}'")
            
            # 执行工作流
            final_state = await self.graph.ainvoke(initial_state, config)
            
            # 构建结果
            result = {
                "success": len(final_state["errors"]) == 0,
                "workflow_id": actual_workflow_id,
                "thread_id": actual_thread_id,
                "query": query,
                "task_type": task_type,
                "results": final_state["results"],
                "errors": final_state["errors"],
                "agent_status": {k: v.value for k, v in final_state["agent_status"].items()},
                "completed_tasks": list(final_state["completed_tasks"]),
                "loop_count": final_state["loop_count"],
                "messages": [msg.content for msg in final_state["messages"]],
                "metadata": final_state["metadata"],
                "start_time": final_state["metadata"].get("workflow_start_time"),
                "end_time": datetime.now().isoformat(),
                "duration_ms": int((datetime.now() - final_state["timestamp"]).total_seconds() * 1000)
            }
            
            logger.info(f"工作流执行完成: workflow_id='{actual_workflow_id}', success={result['success']}, errors={len(result['errors'])}")
            return result
            
        except Exception as e:
            logger.error(f"工作流执行失败: {str(e)}")
            return {
                "success": False,
                "error": str(e),
                "workflow_id": workflow_id or str(uuid.uuid4()),
                "query": query,
                "task_type": task_type,
                "timestamp": datetime.now().isoformat()
            }
    
    async def pause_workflow(self, thread_id: str) -> bool:
        """暂停工作流
        
        Args:
            thread_id: 工作流线程ID
            
        Returns:
            是否成功暂停
        """
        try:
            logger.info(f"暂停工作流: thread_id='{thread_id}'")
            # 这里可以实现工作流暂停逻辑
            # 例如：更新状态到数据库或内存中
            return True
        except Exception as e:
            logger.error(f"暂停工作流失败: {str(e)}")
            return False
    
    async def resume_workflow(self, thread_id: str) -> Dict[str, Any]:
        """恢复工作流
        
        Args:
            thread_id: 工作流线程ID
            
        Returns:
            恢复后的工作流状态
        """
        try:
            logger.info(f"恢复工作流: thread_id='{thread_id}'")
            
            # 获取当前状态
            config = {"configurable": {"thread_id": thread_id}}
            state = self.graph.get_state(config)
            
            if not state:
                raise WorkflowError(f"未找到工作流: {thread_id}")
            
            # 更新状态
            current_state = state.values.copy()
            
            # 将所有暂停的任务状态更新为IDLE
            for agent, status in current_state["agent_status"].items():
                if status == AgentStatus.PAUSED:
                    current_state["agent_status"][agent] = AgentStatus.IDLE
            
            # 继续执行
            final_state = await self.graph.ainvoke(current_state, config)
            
            # 构建结果
            result = {
                "success": True,
                "workflow_id": current_state.get("workflow_id", thread_id),
                "thread_id": thread_id,
                "agent_status": {k: v.value for k, v in final_state["agent_status"].items()},
                "completed_tasks": list(final_state["completed_tasks"]),
                "errors": final_state["errors"]
            }
            
            return result
            
        except Exception as e:
            logger.error(f"恢复工作流失败: {str(e)}")
            raise WorkflowError(f"恢复工作流失败: {str(e)}")
    
    async def cancel_workflow(self, thread_id: str) -> bool:
        """取消工作流
        
        Args:
            thread_id: 工作流线程ID
            
        Returns:
            是否成功取消
        """
        try:
            logger.info(f"取消工作流: thread_id='{thread_id}'")
            # 这里可以实现工作流取消逻辑
            # 例如：清除状态、释放资源等
            return True
        except Exception as e:
            logger.error(f"取消工作流失败: {str(e)}")
            return False
    
    def export_workflow_graph(self, file_path: str = "workflow_graph.png") -> bool:
        """导出工作流图（如果支持可视化）
        
        Args:
            file_path: 导出文件路径
            
        Returns:
            是否成功导出
        """
        try:
            if self.config.enable_visualization and self.graph:
                # 这里可以实现工作流图导出逻辑
                # 例如：使用langgraph的可视化功能
                logger.info(f"导出工作流图到: {file_path}")
                return True
            else:
                logger.warning("工作流可视化未启用")
                return False
        except Exception as e:
            logger.error(f"导出工作流图失败: {str(e)}")
            return False

    def get_workflow_status(self, thread_id: str) -> Dict[str, Any]:
        """获取工作流状态"""
        try:
            config = {"configurable": {"thread_id": thread_id}}
            state = self.graph.get_state(config)
            
            if state:
                return {
                    "thread_id": thread_id,
                    "status": "active",
                    "state": {
                        "messages": len(state.values.get("messages", [])),
                        "errors": len(state.values.get("errors", [])),
                        "agent_status": {k: v.value for k, v in state.values.get("agent_status", {}).items()},
                        "timestamp": state.values.get("timestamp", datetime.now()).isoformat()
                    }
                }
            else:
                return {
                    "thread_id": thread_id,
                    "status": "not_found"
                }
                
        except Exception as e:
            logger.error(f"获取工作流状态失败: {str(e)}")
            return {
                "thread_id": thread_id,
                "status": "error",
                "error": str(e)
            }

    def list_agents(self) -> List[str]:
        """列出已注册的智能体"""
        return list(self.agents.keys())

    def list_tools(self) -> List[str]:
        """列出已注册的工具"""
        return list(self.tools.keys())